package com.example.chatapp.websocket;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class WebSocketSessionManager {
    // sessionId -> session
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    // roomId -> set of sessionIds
    private final Map<String, Set<String>> roomMembers = new ConcurrentHashMap<>();
    // username -> sessionId (assume single session per username for simplicity; extend as needed)
    private final Map<String, String> userSession = new ConcurrentHashMap<>();

    public void add(WebSocketSession session, String username, String roomId) {
        sessions.put(session.getId(), session);
        userSession.put(username, session.getId());
        roomMembers.computeIfAbsent(roomId, k -> ConcurrentHashMap.newKeySet()).add(session.getId());
    }

    public void remove(WebSocketSession session) {
        sessions.remove(session.getId());
        userSession.values().removeIf(sid -> sid.equals(session.getId()));
        roomMembers.values().forEach(set -> set.remove(session.getId()));
    }

    public void broadcastToRoom(String roomId, String message) {
        Set<String> members = roomMembers.get(roomId);
        if (members == null) return;
        TextMessage tm = new TextMessage(message);
        for (String sid : members) {
            WebSocketSession s = sessions.get(sid);
            if (s != null && s.isOpen()) {
                sendMessageSafely(s, tm);
            }
        }
    }

    public void broadcast(String message) {
        TextMessage tm = new TextMessage(message);
        for (WebSocketSession s : sessions.values()) {
            if (s.isOpen()) {
                sendMessageSafely(s, tm);
            }
        }
    }

    public void sendToUser(String username, String message) {
        String sid = userSession.get(username);
        if (sid == null) return;
        WebSocketSession s = sessions.get(sid);
        if (s != null && s.isOpen()) {
            sendMessageSafely(s, new TextMessage(message));
        }
    }

    /**
     * 线程安全的消息发送方法
     * 使用 synchronized 确保同一个 session 不会被并发写入
     */
    private void sendMessageSafely(WebSocketSession session, TextMessage message) {
        try {
            synchronized (session) {
                if (session.isOpen()) {
                    session.sendMessage(message);
                }
            }
        } catch (IOException e) {
            System.err.println("[WebSocketSessionManager] Failed to send message to session " 
                + session.getId() + ": " + e.getMessage());
        } catch (Exception e) {
            System.err.println("[WebSocketSessionManager] Unexpected error sending message: " + e.getMessage());
        }
    }

    public Optional<String> getUsernameBySessionId(String sessionId) {
        return userSession.entrySet().stream().filter(e -> e.getValue().equals(sessionId)).map(Map.Entry::getKey).findFirst();
    }
}
